#!/usr/bin/env python3
# 指定脚本使用Python 3解释器执行

# 导入必要的库
import rclpy  # ROS 2 Python客户端库，用于创建和管理节点
from rclpy.node import Node  # ROS 2节点基类，所有自定义节点都需继承此类
from livox_interfaces.msg import CustomMsg  # Livox激光雷达的自定义点云消息类型
from sensor_msgs.msg import PointCloud2, PointField  # ROS标准点云消息及字段定义
import numpy as np  # 用于高效的数值计算和数组操作
import struct  # 用于二进制数据的打包与解包，构建PointCloud2消息时使用
import open3d as o3d  # 用于3D点云处理、可视化和聚类分析
from threading import Thread, Lock  # 多线程支持，Lock用于线程间同步
import time  # 用于计时和控制程序执行速度
import multiprocessing  # 用于并行计算，提高点云处理效率

# 定义PointCloud2消息的字段结构，描述点云数据的组成
FIELDS = [
    # X坐标：偏移0字节，32位浮点数，数量1
    PointField(name='x', offset=0, datatype=PointField.FLOAT32, count=1),
    # Y坐标：偏移4字节（紧跟X之后），32位浮点数，数量1
    PointField(name='y', offset=4, datatype=PointField.FLOAT32, count=1),
    # Z坐标：偏移8字节（紧跟Y之后），32位浮点数，数量1
    PointField(name='z', offset=8, datatype=PointField.FLOAT32, count=1),
    # 反射率：偏移12字节（紧跟Z之后），8位无符号整数，数量1
    PointField(name='intensity', offset=12, datatype=PointField.UINT8, count=1),
]
POINT_STEP = 13  # 每个点占用的字节数：3个float(4字节×3) + 1个uint8(1字节) = 13字节
ROW_STEP = 0  # 每行的字节数，后续会根据点云数量动态计算


def process_chunk(chunk, distance_threshold, keep_front, area_half_width, area_half_height, 
                 exclude_area_params):
    """
    并行处理点云数据块的过滤函数，实现距离过滤、区域选择和特定区域屏蔽
    
    参数:
        chunk: 点云数据块 (numpy数组，每行代表一个点[x,y,z,intensity])
        distance_threshold: 距离阈值 (米)，只保留该距离内的点
        keep_front: 是否只保留正前方区域的点
        area_half_width: 前方区域的半宽度 (米)
        area_half_height: 前方区域的半高度 (米)
        exclude_area_params: 要排除的区域参数（包含x_min, x_max, y_min, y_max, z_min, z_max）
        
    返回:
        过滤后的点云数据块
    """
    # 1. 距离过滤：计算每个点到原点的平方距离（避免开平方运算提高效率）
    # chunk[:, :3]取每个点的x,y,z坐标，**2计算平方，sum(axis=1)计算每个点的平方和
    squared_distances = np.sum(chunk[:, :3]** 2, axis=1)
    # 创建距离掩码：只保留平方距离小于等于阈值平方的点
    distance_mask = squared_distances <= (distance_threshold **2)
    
    # 2. 前方区域过滤：仅保留传感器正前方特定范围内的点
    if keep_front:
        # 判断点是否在X轴正方向（传感器前方）
        in_front_mask = (chunk[:, 0] > 0)
        # 判断点是否在设定的宽度范围内（Y轴）
        in_width_mask = (np.abs(chunk[:, 1]) <= area_half_width)
        # 判断点是否在设定的高度范围内（Z轴）
        in_height_mask = (np.abs(chunk[:, 2]) <= area_half_height)
        
        # 合并三个条件：同时满足在前方、宽度范围内和高度范围内
        front_area_mask = np.logical_and(np.logical_and(in_front_mask, in_width_mask), in_height_mask)
        # 最终过滤掩码：同时满足距离条件和前方区域条件
        filter_mask = np.logical_and(distance_mask, front_area_mask)
    else:
        # 如果不启用前方区域过滤，则直接使用距离掩码
        filter_mask = distance_mask
    
    # 3. 排除指定区域内的点（例如车辆自身或不需要的区域）
    in_exclude_area = (
        # X轴范围检查
        (chunk[:, 0] >= exclude_area_params['x_min']) &
        (chunk[:, 0] <= exclude_area_params['x_max']) &
        # Y轴范围检查
        (chunk[:, 1] >= exclude_area_params['y_min']) &
        (chunk[:, 1] <= exclude_area_params['y_max']) &
        # Z轴范围检查
        (chunk[:, 2] >= exclude_area_params['z_min']) &
        (chunk[:, 2] <= exclude_area_params['z_max'])
    )
    
    # 最终掩码：保留符合过滤条件且不在排除区域内的点（~表示取反）
    final_mask = np.logical_and(filter_mask, ~in_exclude_area)
    
    # 应用掩码，返回过滤后的点云块
    return chunk[final_mask]


class LidarProcessor(Node):
    """LiDAR点云处理器节点类，实现点云数据的接收、处理、发布和可视化"""
    
    def __init__(self):
        """初始化LiDAR处理器节点，设置订阅者、发布者和各种参数"""
        # 调用父类Node的构造函数，节点名称为"lidar_processor"
        super().__init__('lidar_processor')
        
        # 声明并获取ROS参数：输入点云话题名称，默认值为"/livox/lidar_avia"
        self.declare_parameter('input_cloud', '/livox/lidar_avia')
        # 从参数服务器获取实际的输入话题名称
        input_topic = self.get_parameter('input_cloud').get_parameter_value().string_value

        # 创建订阅者，订阅Livox激光雷达的自定义点云消息
        # 参数分别为：消息类型、话题名称、回调函数、队列长度
        self.sub = self.create_subscription(
            CustomMsg, input_topic, self.callback, 10
        )
        
        # 创建发布者，发布处理后的标准PointCloud2格式点云
        # 参数分别为：消息类型、话题名称、队列长度
        self.pcl_pub = self.create_publisher(
            PointCloud2, '/livox/pointcloud2', 10
        )
        
        # 记录日志信息，显示节点正在订阅和发布的话题
        self.get_logger().info(f"订阅: {input_topic}，发布处理后点云到: /livox/pointcloud2")

        # 初始化Open3D点云对象和可视化窗口
        self.pcd = o3d.geometry.PointCloud()  # 存储点云数据
        self.vis = o3d.visualization.Visualizer()  # 可视化器
        self.vis.create_window()  # 创建可视化窗口
        
        # 设置可视化参数
        self.vis.get_render_option().point_size = 2.0  # 点的显示大小
        self.vis.get_render_option().background_color = np.array([0, 0, 0])  # 黑色背景
        
        # 创建线程锁和标志变量，确保多线程安全
        self.point_cloud_lock = Lock()  # 保护点云数据的线程锁
        self.is_visualizing = True  # 控制可视化线程的运行状态
        self.is_first_frame = True  # 标记是否是第一帧点云数据
        
        # 启动单独的线程用于点云可视化，避免阻塞主处理流程
        self.visualization_thread = Thread(target=self.visualization_loop)
        self.visualization_thread.daemon = True  # 设置为守护线程，主程序退出时自动结束
        self.visualization_thread.start()  # 启动可视化线程
        
        # 点云计数和保存相关参数
        self.frame_count = 0  # 帧计数器，记录处理的点云帧数
        self.save_path = "/path/to/save/pointclouds/"  # 点云保存路径（需根据实际情况修改）
        self.save_enabled = False  # 是否启用点云保存功能
        self.use_fixed_color = True  # 是否使用固定颜色显示点云
        
        # 过滤参数设置
        self.distance_threshold = 50.0  # 距离阈值：只保留50米内的点
        self.keep_front = True         # 是否只保留正前方区域的点
        
        # 前方区域尺寸参数
        self.area_width = 2           # 区域宽度（米）
        self.area_height = 2          # 区域高度（米）
        self.area_half_width = self.area_width / 2  # 半宽度（用于计算）
        self.area_half_height = self.area_height / 2  # 半高度（用于计算）
        
        # 聚类分析参数
        self.cluster_eps = 0.2        # DBSCAN聚类算法的邻域半径（米）
        self.cluster_min_points = 20  # 形成聚类所需的最小点数
        self.enable_clustering = True  # 是否启用聚类分析
        self.voxel_size = 0.05         # 体素降采样的体素大小（米）
        
        # 并行处理设置：使用CPU核心数-1个进程（保留1个核心给系统）
        self.num_processes = max(1, multiprocessing.cpu_count() - 1)
        
        # 定义要排除的区域参数（例如车辆自身底部区域）
        self.exclude_area_params = {
            'x_min': 0.0,    # X轴：前方0米开始
            'x_max': 2.8,    # X轴：前方2.8米结束
            'y_min': -1,     # Y轴：左侧1米
            'y_max': 1.0,    # Y轴：右侧1米
            'z_min': -2,     # Z轴：下方2米
            'z_max': -0.1    # Z轴：下方0.1米
        }
        # 记录日志，显示已启用的区域屏蔽信息
        self.get_logger().info(
            f"已启用区域屏蔽: X∈[{self.exclude_area_params['x_min']},{self.exclude_area_params['x_max']}], "
            f"Y∈[{self.exclude_area_params['y_min']},{self.exclude_area_params['y_max']}], "
            f"Z∈[{self.exclude_area_params['z_min']},{self.exclude_area_params['z_max']}]"
        )

    def callback(self, msg: CustomMsg):
        """
        订阅消息的回调函数，处理接收到的LiDAR点云数据
        参数: msg - 接收到的Livox自定义点云消息
        """
        start_time = time.time()  # 记录处理开始时间，用于计算处理耗时
        
        # 将接收到的点云消息转换为NumPy数组，方便后续处理
        # 数组形状为(n, 4)，n为点的数量，4列分别为x,y,z,反射率
        points_array = np.zeros((len(msg.points), 4), dtype=np.float32)
        for i, p in enumerate(msg.points):
            points_array[i, 0] = p.x  # X坐标（正前方）
            points_array[i, 1] = p.y  # Y坐标（横向）
            points_array[i, 2] = p.z  # Z坐标（上下）
            points_array[i, 3] = p.reflectivity  # 反射率值
        
        # 并行过滤点云：将点云分成多个块，使用多进程同时处理
        # 将点云数组按进程数分割成多个子块
        chunks = np.array_split(points_array, self.num_processes)
        # 创建进程池，使用starmap方法并行调用process_chunk函数
        with multiprocessing.Pool(processes=self.num_processes) as pool:
            # 为每个数据块准备参数，调用process_chunk处理
            results = pool.starmap(
                process_chunk, 
                [
                    (chunk, self.distance_threshold, self.keep_front, 
                     self.area_half_width, self.area_half_height, self.exclude_area_params) 
                    for chunk in chunks
                ]
            )
        # 合并所有处理结果，得到完整的过滤后点云
        filtered_points = np.concatenate(results)
        
        # 如果过滤后的点太少，不进行聚类（优化性能）
        if len(filtered_points) < 50:
            # 每10帧记录一次日志，避免日志过多
            if self.frame_count % 10 == 0:
                self.get_logger().info(f"过滤后点太少（{len(filtered_points)}点），跳过聚类")
            # 所有点都设为白色
            colors = np.ones((len(filtered_points), 3))
        else:
            # 创建Open3D点云对象并进行降采样，减少计算量
            pcd = o3d.geometry.PointCloud()
            # 设置点云坐标（取前3列x,y,z）
            pcd.points = o3d.utility.Vector3dVector(filtered_points[:, :3])
            # 体素降采样：将空间划分为体素网格，每个体素保留一个点
            downsampled_pcd = pcd.voxel_down_sample(voxel_size=self.voxel_size)
            
            # 如果启用聚类分析
            if self.enable_clustering:
                # 使用DBSCAN算法进行聚类
                # 参数：邻域半径、最小点数、是否打印进度
                labels = np.array(downsampled_pcd.cluster_dbscan(
                    eps=self.cluster_eps, 
                    min_points=self.cluster_min_points, 
                    print_progress=False
                ))
                
                # 将降采样点云的聚类标签映射回原始点云
                # 创建KDTree用于快速最近邻搜索
                pcd_tree = o3d.geometry.KDTreeFlann(downsampled_pcd)
                # 初始化原始点云的标签数组，-1表示未分类（噪声）
                full_labels = np.ones(len(filtered_points), dtype=np.int32) * -1
                
                # 对每个原始点，找到最近的降采样点并复制其聚类标签
                for i in range(len(filtered_points)):
                    # 搜索最近的1个点
                    [k, idx, _] = pcd_tree.search_knn_vector_3d(filtered_points[i, :3], 1)
                    full_labels[i] = labels[idx[0]]  # 复制标签
                
                # 生成聚类颜色
                max_label = full_labels.max()  # 获取最大标签值（聚类数量-1）
                self.get_logger().info(f"检测到 {max_label + 1} 个聚类（降采样后点云）")
                
                # 初始化颜色数组，形状为(n, 3)，每行为RGB颜色值
                colors = np.zeros((len(filtered_points), 3))
                # 为每个聚类分配不同的颜色
                for label in range(-1, max_label + 1):
                    label_mask = full_labels == label  # 创建当前标签的掩码
                    if label == -1:
                        # 噪声点设为绿色
                        colors[label_mask] = [0, 1, 0]
                    else:
                        # 为不同聚类分配不同色相的颜色
                        hue = label / (max_label + 1)  # 色相值在0-1之间
                        # 转换HSV到RGB，使用较高的饱和度和明度，使颜色更鲜艳
                        rgb = self.hsv_to_rgb(hue, 0.8, 0.9)
                        colors[label_mask] = rgb
            else:
                # 如果不启用聚类，基于反射率设置颜色
                # 将反射率归一化到0-1范围
                intensities = filtered_points[:, 3] / 255.0
                colors = np.zeros((len(filtered_points), 3))
                # 反射率控制红色通道（反射率越高，红色越浓）
                colors[:, 0] = intensities
        
        # 创建并发布ROS标准PointCloud2消息
        pcl_msg = self.create_pointcloud2_msg(msg.header, filtered_points)
        self.pcl_pub.publish(pcl_msg)
        
        # 更新可视化（使用锁确保线程安全，避免同时读写点云数据）
        with self.point_cloud_lock:
            # 提取点的坐标数据（x,y,z）
            points_np = filtered_points[:, :3]
            
            # 首次接收到点云时，分析坐标范围并可能应用缩放
            if self.is_first_frame:
                # 计算点云坐标的最小值和最大值
                x_min, y_min, z_min = np.min(points_np, axis=0)
                x_max, y_max, z_max = np.max(points_np, axis=0)
                # 记录坐标范围日志
                self.get_logger().info(
                    f"坐标范围: x=[{x_min:.2f}, {x_max:.2f}], "
                    f"y=[{y_min:.2f}, {y_max:.2f}], "
                    f"z=[{z_min:.2f}, {z_max:.2f}]"
                )
                
                # 如果坐标值过大（可能是毫米单位），应用缩放因子转换为米
                scale = 1.0
                max_coord = max(abs(x_max), abs(x_min), abs(y_max), abs(y_min), abs(z_max), abs(z_min))
                if max_coord > 1000:
                    scale = 0.001  # 毫米转米
                    self.get_logger().info(f"应用缩放因子: {scale}")
                    points_np *= scale  # 应用缩放
            
            # 更新Open3D点云数据
            self.pcd.points = o3d.utility.Vector3dVector(points_np)  # 设置坐标
            self.pcd.colors = o3d.utility.Vector3dVector(colors)      # 设置颜色
            
            # 如果启用保存功能，则保存当前帧点云
            if self.save_enabled:
                self.save_current_frame()
            
            # 标记不再是第一帧
            self.is_first_frame = False
        
        # 计算并记录处理时间
        self.frame_count += 1  # 帧计数加1
        # 计算处理耗时（毫秒）
        process_time = (time.time() - start_time) * 1000
        # 记录日志，包含帧号、原始点数、过滤后点数和处理时间
        self.get_logger().info(
            f"帧 {self.frame_count}: 原始{len(msg.points)}点 → 过滤后{len(filtered_points)}点 → "
            f"处理时间{process_time:.2f}ms"
        )

    def hsv_to_rgb(self, h, s, v):
        """
        将HSV颜色空间转换为RGB颜色空间
        参数:
            h: 色相 (0-1)
            s: 饱和度 (0-1)
            v: 明度 (0-1)
        返回:
            RGB颜色值列表 [r, g, b]，每个值在0-1之间
        """
        h6 = h * 6  # 将色相映射到0-6范围
        i = np.floor(h6)  # 整数部分，用于确定颜色区间
        f = h6 - i  # 小数部分，用于插值计算
        p = v * (1 - s)
        q = v * (1 - s * f)
        t = v * (1 - s * (1 - f))
        
        i = int(i) % 6  # 确保i在0-5范围内
        # 根据不同区间返回对应的RGB值
        if i == 0:
            return [v, t, p]
        elif i == 1:
            return [q, v, p]
        elif i == 2:
            return [p, v, t]
        elif i == 3:
            return [p, q, v]
        elif i == 4:
            return [t, p, v]
        else:  # i == 5
            return [v, p, q]

    def create_pointcloud2_msg(self, header, points_array):
        """
        创建ROS标准的PointCloud2消息
        参数:
            header: 消息头，包含时间戳和坐标系信息
            points_array: 点云数据数组，形状为(n, 4)
        返回:
            构建好的PointCloud2消息
        """
        pcl_msg = PointCloud2()
        pcl_msg.header = header  # 设置消息头（继承自输入消息的时间戳和坐标系）
        pcl_msg.height = 1  # 点云为无序点云，高度设为1
        pcl_msg.width = len(points_array)  # 点云宽度为点的数量
        pcl_msg.fields = FIELDS  # 设置点云字段定义
        pcl_msg.is_bigendian = False  # 使用小端字节序
        pcl_msg.point_step = POINT_STEP  # 每个点的字节数
        # 行步长 = 每个点的字节数 × 点的数量
        pcl_msg.row_step = POINT_STEP * pcl_msg.width
        pcl_msg.is_dense = False  # 点云中可能包含无效点（NaN/Inf）
        
        # 将点数据打包为二进制字节流
        data = []
        for point in points_array:
            x, y, z, intensity = point  # 解包点的四个分量
            # 打包为二进制数据：3个float(4字节) + 1个unsigned char(1字节)
            data.append(struct.pack('fffB', x, y, z, int(intensity)))
        # 将所有点的二进制数据连接成一个字节流
        pcl_msg.data = b''.join(data)
        
        return pcl_msg

    def visualization_loop(self):
        """
        可视化线程的主循环，持续更新点云显示
        单独线程运行，避免阻塞点云处理流程
        """
        is_first_visualization = True  # 标记是否是首次可视化
        
        # 只要可视化标志为True，就持续运行
        while self.is_visualizing:
            # 加锁访问点云数据，确保线程安全
            with self.point_cloud_lock:
                # 首次可视化且点云不为空
                if is_first_visualization and len(self.pcd.points) > 0:
                    # 将点云添加到可视化窗口
                    self.vis.add_geometry(self.pcd)
                    
                    # 设置初始视图参数
                    ctr = self.vis.get_view_control()
                    ctr.set_front([0, 0, 1])  # 相机前方方向
                    ctr.set_up([0, 1, 0])     # 相机上方方向
                    ctr.set_lookat([0, 0, 0]) # 相机观察点
                    ctr.set_zoom(0.4)         # 相机缩放级别
                    
                    is_first_visualization = False  # 标记为已完成首次可视化
                elif not is_first_visualization:
                    # 更新点云几何数据
                    self.vis.update_geometry(self.pcd)
                    # 处理窗口事件（如鼠标、键盘输入）
                    self.vis.poll_events()
                    # 更新渲染器，刷新显示
                    self.vis.update_renderer()
            
            # 控制可视化刷新频率（约33Hz）
            time.sleep(0.03)

    def save_current_frame(self):
        """保存当前帧点云到指定路径，格式为PCD文件"""
        # 生成文件名，包含帧号（6位数字，不足补零）
        filename = f"{self.save_path}/pointcloud_{self.frame_count:06d}.pcd"
        # 使用Open3D保存点云
        o3d.io.write_point_cloud(filename, self.pcd)
        # 记录日志，显示保存成功
        self.get_logger().info(f"已保存点云: {filename}")

    def destroy_node(self):
        """节点销毁时的清理工作"""
        self.is_visualizing = False  # 停止可视化线程
        self.vis.destroy_window()    # 关闭可视化窗口
        super().destroy_node()       # 调用父类的销毁方法


def main(args=None):
    """程序主入口函数"""
    rclpy.init(args=args)  # 初始化ROS 2 Python客户端库
    node = LidarProcessor()  # 创建LiDAR处理器节点
    
    try:
        rclpy.spin(node)  # 运行节点，处理回调函数
    except KeyboardInterrupt:
        # 捕获键盘中断（Ctrl+C），优雅退出
        pass
    finally:
        node.destroy_node()  # 销毁节点
        rclpy.shutdown()    # 关闭ROS 2客户端库


# 如果作为主程序运行，则调用main函数
if __name__ == '__main__':
    main()
